Keras BN层 探坑

发表于 | 分类于

问题发生在我使用Keras的Resnet50预训练模型时。我希望用Resnet50作为一个主干网络，通过在后面添加全连接层来实现我的回归任务。训练的时候效果不错，loss经过很多次训练后收敛了，下降到了一个较低的值。但是在测试阶段，预测结果很差，loss很高，即使我使用了训练数据去预测。最后发现问题出现在了BN(Batch Normalization)层。

问题介绍

首先需要说明的是，不同Keras版本的解决方法是不一样的。我的版本如下：

  • tensorflow = 2.2.0

由于tensorflow2.0已经集成了Keras，所以我使用的是tensorflow里的Keras，也就是这样：

1
import tensorflow.keras as keras

网上关于BN层解决方案有些还是tensorflow1.x时代的，因此在尝试的时候应该先确认你的版本。这里有一些我搜索的时候比较经典的几篇博客：

注意他们的版本和时间，离我写这篇博客的时候已经一两年了，至少对于我的版本已经不可用了。

接下来我们简单看一下问题：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.layers import Dense,GlobalAveragePooling2D,Input
from tensorflow.keras import Sequential
# 假设我们有一些图片，我们希望输入这些图片，能够回归到3个参数
# 载入图片
dataset = load_dataset()
x_train,y_train = dataset.x,dataset.y
# 使用预训练的Resnet50
resnet = ResNet50(weights='imagenet', include_top=False,input_shape=(256,256,3))
# 添加全局池化层和全连接层
global_average_layer = GlobalAveragePooling2D()
output_layer = Dense(3)
model = Sequential([resnet,global_average_layer, output_layer])
model.compile(loss='mean_squared_error',optimizer=opt)
model.fit(x_train,y_train,epochs=40)
# 测试
model.evaluate(x_train,y_train)

理论上来说，evaluate出的结果应该和刚刚训练的结果一致。训练结束后，网络的参数都固定了，这时候把刚刚的训练数据放进去测试，结果应该和最后一次训练结果一致，可是结果大相径庭。问题就在于Resnet50中的BN层。

什么是BN层

BN的基本思想：因为深层神经网络在做非线性变换前的激活输入值（就是那个y=wx+b，x是输入）随着网络深度加深或者在训练过程中，其分布逐渐发生偏移或者变动，之所以训练收敛慢，一般是整体分布逐渐往非线性函数的取值区间的上下限两端靠近，所以这导致反向传播时低层神经网络的梯度消失，这是训练深层神经网络收敛越来越慢的本质原因，而BN就是通过一定的规范化手段，把每层神经网络任意神经元这个输入值的分布强行拉回到均值为0方差为1的标准正态分布，其实就是把越来越偏的分布强制拉回比较标准的分布，这样使得激活输入值落在非线性函数对输入比较敏感的区域，这样输入的小变化就会导致损失函数较大的变化，意思是这样让梯度变大，避免梯度消失问题产生，而且梯度变大意味着学习收敛速度快，能大大加快训练速度。

BN在2014年由Loffe和Szegedy提出，它长这个样子：

就是把输入变成一个均值为0，方差为1的正态分布，再进行一个仿射变换(可以想象把正态分布图进行拉伸)。从名字中的Batch就可以看出，这里的均值和方差都是对于一个batch而言的。那么问题来了，训练时候好说，就用当前batch来计算好了，测试的时候怎么办，如果测试的时候只有一个sample呢？，是补成一个batch还是怎么做呢？Keras里是用移动均值和方差，也就是使用历史数据来计算。 所以问题就出在这儿，训练时和测试时，BN层执行的操作是不一样的，因此训练和测试结果不一致。

解决方案

这里不讨论Keras的历史问题，比如某个版本冻结了BN层和没冻结一样等。我们只考虑当前版本如何处理。现在问题就是训练时和测试时执行的操作不一致，那么我们要么让测试去贴合训练，要么就训练阶段贴合测试阶段。

测试端修改

Keras里是有一个变量learning_phase来控制当前是训练还是测试模式的，理论上，我们可以在测试前，强制把模式设置为训练模式，这样测试时不就会按照训练阶段执行了吗？但是我试了没有用，我猜测是不管你外面怎么修改模式，它进行测试的时候都会改成测试模式。果然，在predict的源码里，调用了下面这个函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def predict_step(self, data):
  """The logic for one inference step.
  This method can be overridden to support custom inference logic.
  This method is called by `Model.make_predict_function`.
  This method should contain the mathemetical logic for one step of inference.
  This typically includes the forward pass.
  Configuration details for *how* this logic is run (e.g. `tf.function` and
  `tf.distribute.Strategy` settings), should be left to
  `Model.make_predict_function`, which can also be overridden.
  Arguments:
    data: A nested structure of `Tensor`s.
  Returns:
    The result of one inference step, typically the output of calling the
    `Model` on data.
  """
  data = data_adapter.expand_1d(data)
  x, _, _ = data_adapter.unpack_x_y_sample_weight(data)
  return self(x, training=False)

可以看到，返回的参数training是写死的False，把它改成True问题就迎刃而解。但是直接这样改不是万全之策，极有可能造成其他问题。当然你可以加一个判断，根据当前的learning_phase值来决定training的值。不过直接改包的源码毕竟不太推荐，所以我们可以使用第二种方法：修改训练端。

训练端修改

如果我们让训练阶段执行和测试阶段同样的操作，也是能解决问题的。
TF为后端时，BN有一个参数是training，控制归一化时用的是当前Batch的均值和方差（训练模式）还是移动均值和方差（测试模式）。那我们只要把Resnet的所有BN层的training都修改为测试模式即可。

1
2
3
4
5
6
7
8
# 上面使用resnet
for layer in resnet.layers:
    layerName = str(layer.name)
    if layerName[-2:] == "bn":
        layer.trainable = False
# 下面在resnet后添加其他层

实验一下就会发现，问题完美解决。